import torch
import torch.nn as nn
import math



class Embedding(nn.Module):
    def __init__(self, vocab_size, d_model):
        super(Embedding, self).__init__()
        self.vocab_size = vocab_size
        self.embedding_size = d_model

        self.embedding = nn.Embedding(num_embeddings=vocab_size, embedding_dim=d_model)

    def forward(self, x):
        return self.embedding(x) * math.sqrt(self.embedding_size)

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=60):
        super(PositionalEncoding, self).__init__()

        # 定义dropout层
        self.dropout = nn.Dropout(p=dropout)

        # 思路：位置编码矩阵 + 特征矩阵 相当于给特征增加了位置信息
        # 定义位置编码矩阵PE
        pe = torch.zeros(max_len, d_model)

        # 定义位置列矩阵position -->(max_len,1)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)

        # 定义变化矩阵 得到公式的分母 -->(1,d_model)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() *
                             -(math.log(10000.0) / d_model))

        # 位置列矩阵@变化矩阵 (max_len,1) @ (d_model/2) = (max_len, d_model/2)
        # 矩阵相乘也就是行列对应位置相乘再相加，其含义，给每一个列属性增加位置编码信息
        my_matmulres = position * div_term

        # 给位置编码矩阵奇数列赋值sin
        pe[:, 0::2] = torch.sin(my_matmulres)
        # 偶数列赋值cos
        pe[:, 1::2] = torch.cos(my_matmulres)

        # 升维
        pe = pe.unsqueeze(0)

        # 把pe位置编码矩阵注册成模型的持久缓冲区buffer；模型保存再加载时，可以跟模型一样，一同加载
        # 什么是buffer：对模型效果有帮助的，但是却不是模型结构中超参数或者参数，不参加模型训练
        self.register_buffer('pe', pe)
    def forward(self, x):

        x = x + self.pe[:, :x.size(1)]
        return self.dropout(x)



        
